# -*- coding: utf-8 -*-
#this module has some class about amber.
import re
# a molecular's amber prmtop information.
# read from a amber prmtop file.
#从amber的prmtop文件中读取分子的信息。
import numpy as np
class prmtop():
  elementDict = {'1':'H','6':'C','7':'N','8':'O','11':'NA','15':'P','16':'S','17':'CL'}
  def __init__(self,name):
    self.prm = open(name,'r')
#读某个特定的字段。
  def __field(self):
    result = [];line = self.prm.readline()
    formats = re.findall('\d+',line)
    times = int(formats[0]);length = int(formats[1])
    while 1:
      line = self.prm.readline()
      for m in xrange(times):
        result.append(line[m*length:(m+1)*length])
      if '%FLAG' == line[:5]:
        break
    self.prm.seek(-81,1)
    result = [m.strip() for m in result]
    return result

  def read(self):
    while 1:
      line = self.prm.readline()
      if '%FLAG POINTERS' in line:
        self.prm.readline()
        line = self.prm.readline()
        self.atomNu=int(line[0:8])
        line = self.prm.readline()
        self.residueNu=int(line[8:16])
      elif '%FLAG ATOM_NAME' in line:
        self.atomName = self.__field()[:self.atomNu]
      elif '%FLAG CHARGE' in line:
        self.atomCharge = self.__field()[:self.atomNu]
        self.atomCharge = [float(m) for m in self.atomCharge]
      elif '%FLAG ATOMIC_NUMBER' in line:
        self.atomicNumber = self.__field()[:self.atomNu]
        self.element = [self.__class__.elementDict[m] for m in self.atomicNumber]
      elif '%FLAG RESIDUE_LABEL' in line:
        self.residueName = self.__field()[:self.residueNu]
      elif '%FLAG RESIDUE_POINTER' in line:
        pointer = self.__field()[:self.residueNu]
        pointer = [int(m) for m in pointer]
        self.residuePointer = []
        for i in xrange(len(pointer)-1):
          self.residuePointer.append((pointer[i],pointer[i+1]-1))
        self.residuePointer.append((pointer[i+1],self.atomNu))
        
      if line == '':
        break     
#读取amber输出的某一帧轨迹，这个是轨迹因为是改amber的源代码输出的，所以只有体系
#所有原子的坐标,直接全部读取即可。
def readAmberCrd(name):
  crd0 = open(name,'r')
  crd = crd0.read()
  crd0.close()
  crd = np.fromstring(crd,sep=' ')
  crd = crd.reshape(crd.size/3,3)
  return crd

def readAmberInpcrd(name):
  inpcrd0 = open(name,'r')
  inpcrd0.readline();inpcrd0.readline()
  inpcrd = inpcrd0.readlines()
  inpcrd0.close()
  inpcrd = ''.join(inpcrd[:-1])
  inpcrd = np.fromstring(inpcrd,sep=' ')
  inpcrd = inpcrd.reshape(inpcrd.size/3,3)
  return inpcrd
  
if __name__ == '__main__':
  def testPrm():
    prm = prmtop('1emv1.prmtop')
    prm.read()
    print prm.atomName[:10]
    print prm.atomCharge[:10]
    print prm.residueName[:10]
    print prm.residuePointer[:10]
    print prm.atomName[-10:]
    print prm.atomCharge[-10:]
    print prm.residueName[-10:]
    print prm.residuePointer[-10:]
    print prm.atomicNumber[:10]
    print prm.atomicNumber[-10:]
    print prm.element[:10]
    print prm.element[-10:]
    print prm.atomNu,prm.residueNu
  testPrm()

#  crd = readAmberCrd('1.mdcrd')
#  print crd[:10]
#  crd = readAmberInpcrd('1emv.inpcrd')
#  crd = readAmberCrd('1emv.inpcrd_clean')
#  print crd[:10]
#  print crd[-10:]
